import pybullet as p
import pybullet_data
import time

p.connect(p.GUI)
p.setGravity(0, 0, -9.8)
p.setAdditionalSearchPath(pybullet_data.getDataPath()) # 添加pybullet_data路径

# 加载cartpole模型，它已经包含了一个旋转关节（第二个关节，索引1）
cartpole_id = p.loadURDF("cartpole.urdf")

# 获取关节信息，确认旋转关节的索引
num_joints = p.getNumJoints(cartpole_id)
for i in range(num_joints):
    joint_info = p.getJointInfo(cartpole_id, i)
    print(f"Joint {i}: Name: {joint_info[1].decode('utf-8')}, Type: {joint_info[2]}")

# 假设旋转关节的索引是1
revolute_joint_index = 1

# 使用速度控制让杆子持续旋转
max_force = 500
target_velocity = 5  # 弧度/秒
p.setJointMotorControl2(bodyUniqueId=cartpole_id,
                       jointIndex=revolute_joint_index,
                       controlMode=p.VELOCITY_CONTROL,
                       targetVelocity=target_velocity,
                       force=max_force)

# 仿真循环
try:
    for _ in range(10000):
        p.stepSimulation()
        # 获取当前关节状态（位置、速度等）
        joint_state = p.getJointState(cartpole_id, revolute_joint_index)
        print(f"Joint Position: {joint_state[0]:.2f}, Velocity: {joint_state[1]:.2f}")
        time.sleep(1./240.)
except KeyboardInterrupt:
    pass

p.disconnect()